#!/usr/bin/python
# -*- coding: UTF-8 -*-
# author: ziping.yu
# Date: 2023.5.10
import rospy
import PyKDL
from geometry_msgs.msg import Twist, Point, PoseStamped
from std_msgs.msg import String
import numpy as np
import math
import sys

class NavigationPlan:
  def __init__(self, vehicle_type, vehicle_id, vehicle_num=0, drone_num=0):
    self.vehicle_type = vehicle_type
    self.vehicle_id = int(vehicle_id)
    self.vehicle_num = int(vehicle_num)
    self.drone_num = int(drone_num)
    self.bias = np.array([[1,0,0], [-1,-1,0], [-1,1,0], [4,3,0], [2,4,0], [2,2,0], [4,-3,0], [2,-4,0], [2,-2,0]]) # only for ugv which in enu coordinate
    self.waypoints = np.array([[[15.1, -0.1],[25.2, -0.2],[30.0,-11.2],[43.7, -0.1],[34.6, 10.6]],
                               [[15.1, -0.1],[25.2, -0.2],[30.0,-11.2],[43.7, -0.1],[34.6, 10.6]],
                               [[15.1, -0.1],[25.2, -0.2],[30.0,-11.2],[43.7, -0.1],[34.6, 10.6]],
                               [[6.9, 14.2],[26.1, 18.5],[37.9, 32.3],[15.3, 41.6],[-7.3, 26.5]], 
                               [[6.9, 14.2],[26.1, 18.5],[37.9, 32.3],[15.3, 41.6],[-7.3, 26.5]],
                               [[6.9, 14.2],[26.1, 18.5],[37.9, 32.3],[15.3, 41.6],[-7.3, 26.5]],
                               [[5.2,-13.2],[-5.5,-29.3],[ 5.3,-39.3],[24.6,-37.2],[21.7,-16.5]], 
                               [[5.2,-13.2],[-5.5,-29.3],[ 5.3,-39.3],[24.6,-37.2],[21.7,-16.5]],  
                               [[5.2,-13.2],[-5.5,-29.3],[ 5.3,-39.3],[24.6,-37.2],[21.7,-16.5]],
                               ])
    '''
    ros parameter configuration
    '''
    self.rate = 30 # pin lv
    
    '''
    control parameter configuration
    '''
    self.kp = 0.5
    self.kv = 1.0
    self.ky = 2.0
    
    '''
    boundary value configuration
    '''
    self.vxy_max = 0.5
    self.angz_max = 2.0
    
    '''
    variable of vehicle
    '''
    self.current_pose = PoseStamped()
    self.current_position = Point()
    self.current_yaw = 0.0
    self.current_vel = Twist()
    self.target_motion = Twist()
    self.target_position = Twist()
    self.target_yaw = 0.0
    self.last_yaw = 0.0
    self.cmd = ''
    self.other_drone_pose = [PoseStamped() for i in range(self.drone_num)]
    self.other_vehicle_pose = [PoseStamped() for i in range(self.vehicle_num)]
    
    '''
    FSM variable
    '''
    self.FSM_EXEC_STATE = {0:"TAKEOFF", 1:"INIT", 2:"GEN_TRAJ", 3:"EXEC"}
    self.exec_state = 0
    self.change_task_flag = False
    self.enable_receive_goal = False
    self.count = 0
    self.is_arrive = False
    self.arrive_count = 0
    self.next_node = 0
    '''
    obstacle avoid
    '''
    self.avoid_finish_flag = True
    self.kp_avoid = 3.0
    self.safe_distance = 0.5 # distance between vehicles for avoidance
    
    rospy.init_node(f"{vehicle_type}_{str(vehicle_id)}_plan")
    '''
    ROS Subscriber
    '''
    self.waypoint_sub = rospy.Subscriber(f"{self.vehicle_type}_{self.vehicle_id}/move_base_simple/goal", 
                                     PoseStamped, self.waypoint_callback)
    self.local_pose_sub = rospy.Subscriber(f"{self.vehicle_type}_{str(self.vehicle_id)}/global_position/pose",
                                           PoseStamped, self.local_pose_callback)
    self.other_drone_pose_sub = [None] * self.drone_num
    self.other_vehicle_pose_sub = [None] * self.vehicle_num
    '''
    ROS Publisher
    '''
    self.vel_flu_pub = rospy.Publisher(f"{self.vehicle_type}_{str(self.vehicle_id)}/cmd_vel_flu", Twist, queue_size=1)
    self.cmd_pub = rospy.Publisher(f"{self.vehicle_type}_{str(self.vehicle_id)}/cmd", String, queue_size=10)
  
  def waypoint_callback(self, msg):
    print("trigger!!!")
    self.target_position = msg.pose.position
    x = msg.pose.orientation.x
    y = msg.pose.orientation.y
    z = msg.pose.orientation.z
    w = msg.pose.orientation.w
    rot = PyKDL.Rotation.Quaternion(x, y, z, w)
    
    res = rot.GetRPY()[2]
    while res > math.pi:
      res -= 2.0 * math.pi
    while res < -math.pi:
      res += 2.0 * math.pi
    if res < 0:
      res = res  + 2.0 * math.pi
    self.target_yaw = res
  
  def local_pose_callback(self, msg):
    self.current_pose = msg
    self.current_position = self.current_pose.pose.position
    self.current_position.x = self.current_position.x + self.bias[self.vehicle_id][0]
    self.current_position.y = self.current_position.y + self.bias[self.vehicle_id][1]
    self.current_position.z = self.current_position.z
    x = self.current_pose.pose.orientation.x
    y = self.current_pose.pose.orientation.y
    z = self.current_pose.pose.orientation.z
    w = self.current_pose.pose.orientation.w
    rot = PyKDL.Rotation.Quaternion(x, y, z, w)
    
    res = rot.GetRPY()[2]
    while res > math.pi:
      res -= 2.0 * math.pi
    while res < -math.pi:
      res += 2.0 * math.pi
    if res < 0:
      res = res  + 2.0 * math.pi
    self.current_yaw = res   # [0, 2pi]
  
  def other_drone_pose_callback(self, msg, id):
    self.other_drone_pose[id] = msg
    self.other_drone_pose[id].pose.position.x = self.other_drone_pose[id].pose.position.x + self.bias[id][0]
    self.other_drone_pose[id].pose.position.y = self.other_drone_pose[id].pose.position.y + self.bias[id][1]
    self.other_drone_pose[id].pose.position.z = self.other_drone_pose[id].pose.position.z
  
  def other_vehicle_pose_callback(self, msg, id):
    self.other_vehicle_pose[id] = msg
    
  def run(self):
    rate = rospy.Rate(self.rate)
    for i in range(self.drone_num):
      self.other_drone_pose_sub[i] = rospy.Subscriber(f"drone_{str(i)}/mavros/local_position/pose",
                                                      PoseStamped, self.other_drone_pose_callback, i)
    for i in range(self.vehicle_num):
      if not i == self.vehicle_id:
        self.other_vehicle_pose_sub[i] = rospy.Subscriber(f"{self.vehicle_type}_{str(i)}/global_position/pose",
                                                          PoseStamped, self.other_vehicle_pose_callback, i)
    while not rospy.is_shutdown():
      self.count += 1
      self.cmd = ''
      if self.exec_state == 0:
        self.target_position.linear.x = self.current_position.x
        self.target_position.linear.y = self.current_position.y
        self.target_motion.angular.z = self.target_motion.angular.z
      
      # move to the initial point
      elif self.exec_state == 1 and self.change_task_flag:
        self.change_task_flag = False
        self.enable_receive_goal = True
        self.change_task_flag = False
        self.init_point()
        # self.Obstable_vehicle_object()  # add obs

      # move in the way
      if self.exec_state == 2 and (self.change_task_flag or self.is_arrive):
          print(self.vehicle_type + '_' + str(self.vehicle_id) + "circle1  flag")
          self.circle1_plan()
          self.change_task_flag = False
      
      distance_tar_cur = self.VectNorm2(self.target_position.linear, self.current_position)
      if distance_tar_cur < 1:
        self.arrive_count += 1
        if self.arrive_count > 3:
          self.is_arrive = True
          self.arrive_count = 0
        else:
          self.is_arrive = False
      else:
        self.arrive_count = 0
        self.is_arrive = False
        
      # task change
      if (self.exec_state == 0) and self.is_arrive and self.count > 100:
        self.change_task_flag = True
        self.exec_state = 1
        self.is_arrive = False
      elif (self.exec_state == 1) and self.is_arrive and self.avoid_finish_flag:
        self.exec_state = 2
        self.is_arrive = False
        self.change_task_flag = True
      
      self.get_control_vel()
      self.vehicle_avoid()
      # print(self.target_motion.linear)
      self.vel_flu_pub.publish(self.target_motion)
      self.cmd_pub.publish(self.cmd)
      
      rate.sleep()
    
  def get_control_vel(self): # PID control, flu coordinate
    dis_curtar = self.VectNorm2(self.target_position.linear, self.current_position)
    vel_total = self.kp * dis_curtar  # velocity
    if vel_total > self.vxy_max:
        vel_total = self.vxy_max
    self.target_yaw = self.pos2ang(self.target_position.linear.x, self.target_position.linear.y, self.current_position.x, self.current_position.y)
    mid_yaw = self.target_yaw - self.current_yaw 
    if mid_yaw > math.pi:
        mid_yaw = mid_yaw - 2 * math.pi
    elif mid_yaw < -math.pi:
        mid_yaw = 2 * math.pi + mid_yaw
    self.target_motion.angular.z = self.ky * mid_yaw
    if self.target_motion.angular.z > self.angz_max:
        self.target_motion.angular.z = self.angz_max
    elif self.target_motion.angular.z < -self.angz_max:
        self.target_motion.angular.z = -self.angz_max
    self.target_motion.linear.x  = vel_total * math.cos(mid_yaw)
    self.target_motion.linear.y = vel_total * math.sin(mid_yaw)
    self.target_motion.linear.z = 0
    # self.target_motion.linear.y = vel_total * math.sin(mid_yaw)
  
  def vehicle_avoid(self):
    self.avoid = []
    for i in range(self.drone_num):
      if not i == self.vehicle_id:
        dis_partner = math.sqrt(
          (self.current_position.x - self.other_vehicle_pose[i].pose.position.x) ** 2 + (self.current_position.y - self.other_vehicle_pose[i].pose.position.y) ** 2
        )
        if dis_partner < self.safe_distance:
          self.avoid.append(i)
    avoid_num = len(self.avoid)
    if avoid_num > 0:
      self.target_motion.linear.x = 0
      self.target_motion.angular.z = 0
      for i in self.avoid:
        print(f"{self.vehicle_id} need avoid!!!")
        # self.target_position.linear.x += self.kp_avoid * (self.current_position.x - self.other_vehicle_pose[i].pose.position.x)
        # self.target_position.linear.y += self.kp_avoid * (self.current_position.y - self.other_vehicle_pose[i].pose.position.y)
        # self.target_position.linear.x = self.kp_avoid * (self.current_position.x - self.other_vehicle_pose[i].pose.position.x)
        # self.target_position.linear.y = self.kp_avoid * (self.current_position.y - self.other_vehicle_pose[i].pose.position.y)
        vec2obj = self.VectDiff(self.other_vehicle_pose[i].pose.position, self.current_position) 
        vec2tar = self.VectDiff(self.target_position.linear, self.current_position)
        diff_ang = self.VectAngDiff(vec2obj, vec2tar)
        if diff_ang > math.pi:
          self.target_motion.linear.x += 0.5
        elif math.fabs(diff_ang) < 0.01:
          self.target_motion.linear.x = self.vxy_max
        else:
          self.target_motion.linear.x -= 0.5
        diff_yaw = self.pos2ang(self.other_vehicle_pose[i].pose.position.x, self.other_vehicle_pose[i].pose.position.y, self.current_position.x, self.current_position.y)
        mid_yaw = diff_yaw - self.current_yaw
        if mid_yaw > math.pi:
          mid_yaw = mid_yaw - 2 * math.pi
        elif mid_yaw < -math.pi:
          mid_yaw = 2 * math.pi + mid_yaw
        self.target_motion.angular.z += -self.ky * mid_yaw
        # TODO: self.target_motion.angular.z
    
  def init_point(self):
    self.target_position.linear.x = self.waypoints[self.vehicle_id][0][0] + self.bias[self.vehicle_id][0]
    self.target_position.linear.y = self.waypoints[self.vehicle_id][0][1] + self.bias[self.vehicle_id][1]
    self.circle1_flag = True
    print(f"[INIT] Toward target point x:{self.target_position.linear.x}, y:{self.target_position.linear.y}")

  def circle1_plan(self):
    self.outer_flag = False
    # self.velxy_max = 2
    size1 = len(self.waypoints[self.vehicle_id])
    # di yi ci jin ru sou suo mo shi xu yao pan duan xia yi ge mu biao dian de suo yin zhi
    if self.change_task_flag:
      self.next_node = 0
      # self.situation2_flag.situation2_id = 1 * 100 + self.vehicle_id   # *Su
      n = 0
      for i in self.waypoints[self.vehicle_id]:
        n = n + 1
        if abs(self.current_position.x - i[0]) < 1:
          self.next_node = n
          break
        elif abs(self.current_position.y - i[1]) < 1:
          if  n == 1:
            self.next_node = 0
          else:
            self.next_node = n
          break
      if self.next_node == size1:
          self.next_node = 0
      self.target_position.linear.x = self.waypoints[self.vehicle_id][self.next_node][0] + self.bias[self.vehicle_id][0]
      self.target_position.linear.y = self.waypoints[self.vehicle_id][self.next_node][1] + self.bias[self.vehicle_id][1]
      self.change_task_flag = False
      # self.situation2_id_pub.publish(self.situation2_flag)
    elif self.is_arrive:
      self.target_position.linear.x = self.waypoints[self.vehicle_id][self.next_node][0] + self.bias[self.vehicle_id][0]
      self.target_position.linear.y = self.waypoints[self.vehicle_id][self.next_node][1] + self.bias[self.vehicle_id][1]
      self.next_node += 1
      if self.next_node == size1:
          self.next_node = 0
  
  def VectNorm3(self, pt1, pt2):
    norm = math.sqrt(pow(pt1.x - pt2.x, 2) + pow(pt1.y - pt2.y, 2) + pow(pt1.z - pt2.z, 2))
    return norm
  
  def VectNorm2(self, pt1, pt2):
    norm = math.sqrt(pow(pt1.x - pt2.x, 2) + pow(pt1.y - pt2.y, 2))
    return norm
  
  def VectDiff(self, endPt, startPt):
        temp = Point()
        temp.x = endPt.x - startPt.x
        temp.y = endPt.y - startPt.y
        return temp
      
  def VectAngDiff(self, v1, v2):
    if isinstance(v1, Point):
      v1 = np.array([v1.x, v1.y])
    if isinstance(v2, Point):
      v2 = np.array([v2.x, v2.y])
    norm_v1 = np.linalg.norm(v1)
    norm_v2 = np.linalg.norm(v2)
    if norm_v1 < 1e-8 or norm_v2 < 1e-8:
        # 向量长度过小，无法计算夹角，返回0
        return 0
    else:
        cos_angle = np.dot(v1, v2) / (norm_v1 * norm_v2)
        angle = np.arccos(np.clip(cos_angle, -1, 1))
        return angle
      
  def pos2ang(self, xa, ya, xb, yb):
    if not xa-xb == 0:
      angle = math.atan2((ya - yb),(xa - xb))
      if (ya-yb > 0) and (angle < 0):  
          angle = angle + math.pi       
      elif (ya-yb < 0) and (angle > 0): 
          angle = angle - math.pi       
      elif ya-yb == 0:
          if xa-xb > 0:
              angle = 0.0
          else:
              angle = math.pi
    else:
      if ya-yb > 0:
          angle = math.pi / 2
      elif ya-yb <0:
          angle = -math.pi / 2
      else:
          angle = 0.0
    if angle < 0:
      angle = angle + 2 * math.pi   # 0 to 2pi
    return angle
  
if __name__ == '__main__':
  planner = NavigationPlan(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
  planner.run()